import numpy as np
from typing import Tuple, List

class CocoBenchmark:
    """COCO基准测试函数集合"""
    
    @staticmethod
    def sphere(x: np.ndarray) -> float:
        """F1: Sphere function
        f(x) = sum(x_i^2)
        """
        return np.sum(x ** 2)
    
    @staticmethod
    def rastrigin(x: np.ndarray) -> float:
        """F3: Rastrigin function
        f(x) = 10n + sum(x_i^2 - 10cos(2πx_i))
        """
        n = len(x)
        return 10 * n + np.sum(x ** 2 - 10 * np.cos(2 * np.pi * x))
    
    @staticmethod
    def rosenbrock(x: np.ndarray) -> float:
        """F8: Rosenbrock function
        f(x) = sum(100(x_{i+1} - x_i^2)^2 + (x_i - 1)^2)
        """
        return np.sum(100 * (x[1:] - x[:-1]**2)**2 + (x[:-1] - 1)**2)
    
    @staticmethod
    def ellipsoidal(x: np.ndarray) -> float:
        """F10: Ellipsoidal function
        f(x) = sum(10^((i-1)/(n-1)) * x_i^2)
        """
        n = len(x)
        i = np.arange(n)
        return np.sum(10**((i-1)/(n-1)) * x**2)
    
    @staticmethod
    def weierstrass(x: np.ndarray) -> float:
        """F16: Weierstrass function
        f(x) = sum(sum(a^k * cos(2πb^k(x_i + 0.5))) - n*sum(a^k * cos(πb^k))
        """
        a = 0.5
        b = 3
        kmax = 20
        n = len(x)
        
        k = np.arange(kmax + 1)
        a_k = a ** k
        b_k = b ** k
        
        term1 = np.sum(a_k[:, np.newaxis] * np.cos(2 * np.pi * b_k[:, np.newaxis] * (x + 0.5)))
        term2 = n * np.sum(a_k * np.cos(np.pi * b_k))
        
        return term1 - term2
    
    @staticmethod
    def schwefel(x: np.ndarray) -> float:
        """F20: Schwefel function
        f(x) = 418.9829n - sum(x_i * sin(sqrt(|x_i|)))
        """
        n = len(x)
        return 418.9829 * n - np.sum(x * np.sin(np.sqrt(np.abs(x))))

    @staticmethod
    def get_problem_info(dim: int) -> Tuple[List[float], List[float], List[float]]:
        """获取问题的边界信息
        
        Args:
            dim: 问题维度
            
        Returns:
            Tuple[List[float], List[float], List[float]]: 下界、上界和最优解
        """
        # 对于大多数COCO函数，搜索空间是[-5, 5]^n
        lower_bounds = [-5.0] * dim
        upper_bounds = [5.0] * dim
        optimal_solution = [0.0] * dim  # 大多数函数在原点处取得最优值
        
        return lower_bounds, upper_bounds, optimal_solution

    @staticmethod
    def evaluate(x: np.ndarray, problem_name: str) -> float:
        """评估指定问题的适应度值
        
        Args:
            x: 决策变量
            problem_name: 问题名称
            
        Returns:
            float: 适应度值
        """
        problem_map = {
            'sphere': CocoBenchmark.sphere,
            'rastrigin': CocoBenchmark.rastrigin,
            'rosenbrock': CocoBenchmark.rosenbrock,
            'ellipsoidal': CocoBenchmark.ellipsoidal,
            'weierstrass': CocoBenchmark.weierstrass,
            'schwefel': CocoBenchmark.schwefel
        }
        
        if problem_name not in problem_map:
            raise ValueError(f"未知的问题名称: {problem_name}")
            
        return problem_map[problem_name](x) 